#include <experimental/filesystem>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <thread>

#include "target_storage.h"

using namespace std;
namespace fs = std::experimental::filesystem;

namespace {
    // const static std::vector<std::string> kLabels = { "person", "bicycle", "car", "motorbike",
    //   "aeroplane","bus", "train", "truck", "boat",
    //   "traffic light", "fire hydrant", "stop sign", "parking meter",
    //   "bench", "bird", "cat", "dog", "horse",
    //   "sheep", "cow", "elephant", "bear", "zebra",
    //   "giraffe", "backpack", "umbrella", "handbag","tie",
    //   "suitcase", "frisbee", "skis", "snowboard", "sports ball",
    //   "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
    //   "tennis racket", "bottle", "wine glass", "cup",
    //   "fork", "knife", "spoon", "bowl", "banana",
    //   "apple", "sandwich", "orange", "broccoli", "carrot",
    //   "hot dog", "pizza", "donut", "cake", "chair",
    //   "sofa", "potted plant", "bed", "dining table", "toilet",
    //   "TV monitor", "laptop", "mouse", "remote", "keyboard",
    //   "cell phone", "microwave", "oven", "toaster", "sink",
    //   "refrigerator", "book", "clock", "vase","scissors",
    //   "teddy bear", "hair drier", "toothbrush"};
    const static std::vector<std::string> kLabels = { "car", "van", "bus", "others" };

    const int kStorageDataQueueSize = 256;

    const int kVideoFPS = 30;

    const std::string DEFAULT_OUTPUT_DIR = "./output";
    const std::string VIDEO_NAME = "compressed.flv";
    const std::string TARGETS_DIR_NAME = "targets";
    const std::string TRAJ_TXT_NAME = "trajectories.txt";
    const std::string SNAPSHOTS_DIR_NAME = "snapshots";
    const std::string IMAGE_DIR_NAME = "images";
    const std::string DETECTIONS_TXT_DIR = "det_results";
}

TargetStorage::TargetStorage() : StorageDataQueue_(kStorageDataQueueSize) , inited_(false) {
}

TargetStorage::~TargetStorage() {
    DestroyResource();
}

void TargetStorage::DestroyResource() {
    while(StorageDataQueue_.size());

    encoder_.DestroyResource();
    dvpp_.DestroyResource();
    AtlasError ret;
    if (stream_ != nullptr) {
        ret = aclrtDestroyStream(stream_);
        if (ret != ACL_ERROR_NONE) {
            ERROR_LOG("destroy stream failed");
        }
        stream_ = nullptr;
    }
    INFO_LOG("end to destroy stream");
}

void TargetStorage::init(uint32_t width, uint32_t height, std::string outputPath, aclrtContext& context) {
    if (inited_) {
        return;
    }
    if (outputPath.empty()) {
        outputPath_ = DEFAULT_OUTPUT_DIR;
    } else {
        outputPath_ = outputPath;
    }
    
    fs::create_directories(outputPath_);
    writer.open((fs::path(outputPath_) / VIDEO_NAME).string(),
                cv::VideoWriter::fourcc('a', 'v', 'c', '1'),
                kVideoFPS, cv::Size(width, height));
    if (!writer.isOpened()) {
        throw std::runtime_error("Cannot open cv::VideoWriter");
    }

    encoder_.InitResource(width, height, outputPath_);

    context_ = context;
    // create stream
    AtlasError ret = aclrtCreateStream(&stream_);
    if (ret != ACL_ERROR_NONE) {
        ERROR_LOG("acl create stream failed");
        return;
    }
    dvpp_.InitResource(stream_);

    thread sender_thread(TargetStorageThreadFunction, (void*)this);
    sender_thread.detach();
    inited_ = true;
}

void TargetStorage::TargetStorageThreadFunction(void* senderSelf) {

    TargetStorage* thisPtr = (TargetStorage*)senderSelf;
    thisPtr->SetAclContext();
    StorageData data;
    while (true) {
        data = thisPtr->StorageDataQueue_.PopAndWait();
        // struct timeval begin;
        // struct timeval end;
        // gettimeofday(&begin, NULL);
        // thisPtr->dvpp_write_image(data.bgr, data.image_id);   // dvpp保存图片有bug，在保持到1002张图片时，读图线程就会一直读图失败，目前不清楚是什么原因
        thisPtr->dvpp_write_video(data.bgr);
        // gettimeofday(&end, NULL);
        // float time_cost = 1000 * (end.tv_sec - begin.tv_sec) + (end.tv_usec - begin.tv_usec) / 1000.000;
        // INFO_LOG("Storage time: %fms", time_cost);
    }
}

Result TargetStorage::StorageDataEnQueue(cv::Mat& bgr, int image_id) {
    StorageData data;
    data = {bgr, image_id-1};

    if (StorageDataQueue_.Push(data)) {
        return SUCCESS;
    }

    ERROR_LOG("storage %d lost data for queue full", 
	                image_id);

    return FAILED;
}

void TargetStorage::cv_write_image(cv::Mat& img, uint32_t frame) {
    auto dir_path = fs::path(outputPath_) / IMAGE_DIR_NAME;
    fs::create_directories(dir_path);
    dir_path /= "/";
    string filename = dir_path.c_str() + to_string(frame) + ".jpg";
    // cout <<filename<<endl;
    cv::imwrite(filename, img);
}

void TargetStorage::dvpp_write_image(const cv::Mat& bgr, uint32_t frame) {

    cv::Mat yuv;
    Utils::BGR2YUV_nv12(bgr, yuv);
    
    ImageData imageData, jpgImage;
    Utils::CopyImageDataToDvpp(imageData, yuv);
    imageData.format = PIXEL_FORMAT_YUV_SEMIPLANAR_420;
    imageData.width = bgr.cols;
    imageData.height = bgr.rows;

    Result ret = (Result)dvpp_.JpegE(jpgImage, imageData);
    if (ret == FAILED) {
        ERROR_LOG("jpg encode failed");
        return;
    }

    auto dir_path = fs::path(outputPath_) / IMAGE_DIR_NAME;
    fs::create_directories(dir_path);
    dir_path /= "/";
    string filename = dir_path.c_str() + to_string(frame) + ".jpg";
    // cout <<filename<<endl;
    FILE * outFileFp = fopen((filename).c_str(), "wb+");
    Utils::WriteToFile(outFileFp, jpgImage.data.get(), jpgImage.size);
}

void TargetStorage::dvpp_encode_image(const cv::Mat& src, ImageData& dst) {

    cv::Mat yuv;
    Utils::BGR2YUV_nv12(src, yuv);
    
    ImageData imageData;
    Utils::CopyImageDataToDvpp(imageData, yuv);
    imageData.format = PIXEL_FORMAT_YUV_SEMIPLANAR_420;
    imageData.width = src.cols;
    imageData.height = src.rows;

    Result ret = (Result)dvpp_.JpegE(dst, imageData);
    if (ret == FAILED) {
        ERROR_LOG("jpg encode failed");
        return;
    }
}

void TargetStorage::cv_write_video(cv::Mat& img) {
    writer.write(img);
}

void TargetStorage::dvpp_write_video(cv::Mat& img) {
    encoder_.DoVencProcess(img);
}

void TargetStorage::save_as_txt(const vector<BBox> &dets, string fileName, uint32_t origin_width, uint32_t origin_height, bool isSaveConf) {
    auto dir_path = fs::path(outputPath_) / DETECTIONS_TXT_DIR;
    fs::create_directories(dir_path);
    dir_path /= "/";
    // imagename.jpg -> imagename.txt
    int fileNameIndex = fileName.find_last_of('/');
    ofstream fp(dir_path / (string(fileName.begin()+fileNameIndex+1, fileName.end()-4) + ".txt"), ios::app);
    for (auto& det : dets) {
        cv::Rect rect = Utils::xywh_cxcywh(det.rect);
        fp  << det.cls_id;
        if (isSaveConf) {
            fp  << " " << det.score;
        }
        fp  << " " << (float)rect.x / origin_width
            << " " << (float)rect.y / origin_height
            << " " << (float)rect.width / origin_width
            << " " << (float)rect.height / origin_height
            << " " << endl;
    }
}


AtlasError TargetStorage::SetAclContext() {
    if (context_ == nullptr) {
        ATLAS_LOG_ERROR("TargetStorage context is null");
        return ATLAS_ERROR_SET_ACL_CONTEXT;
    }
    
    aclError ret = aclrtSetCurrentContext(context_);
    if (ret != ACL_ERROR_NONE) {
        ATLAS_LOG_ERROR("TargetStorage set context failed, error: %d", ret);
        return ATLAS_ERROR_SET_ACL_CONTEXT;
    }
    
    return ATLAS_OK;   
}